{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eec680f5",
   "metadata": {},
   "source": [
    "# LangGraph 단계별 스트리밍 출력\n",
    "\n",
    "이번에는 LangGrpah 의 `stream()` 출력 함수에 대한 조금 더 자세한 설명을 진행합니다.\n",
    "\n",
    "LangGraph 의 스트리밍 출력 함수는 그래프의 각 단계를 스트리밍하는 기능을 제공합니다.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2184f4",
   "metadata": {},
   "source": [
    "참고: 아래의 LangGraph 예제는 이전 섹션의 예제와 동일합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de9d9d8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b5c6228",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Modules\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51549b1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, List, Dict\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langchain_core.tools import tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "from langchain_teddynote.graphs import visualize_graph\n",
    "from langchain_teddynote.tools import GoogleNews\n",
    "\n",
    "\n",
    "########## 1. 상태 정의 ##########\n",
    "# 상태 정의\n",
    "class State(TypedDict):\n",
    "    # 메시지 목록 주석 추가\n",
    "    messages: Annotated[list, add_messages]\n",
    "    dummy_data: Annotated[str, \"dummy\"]\n",
    "\n",
    "\n",
    "########## 2. 도구 정의 및 바인딩 ##########\n",
    "# 도구 초기화\n",
    "# 키워드로 뉴스 검색하는 도구 생성\n",
    "news_tool = GoogleNews()\n",
    "\n",
    "\n",
    "@tool\n",
    "def search_keyword(query: str) -> List[Dict[str, str]]:\n",
    "    \"\"\"Look up news by keyword\"\"\"\n",
    "    news_tool = GoogleNews()\n",
    "    return news_tool.search_by_keyword(query, k=5)\n",
    "\n",
    "\n",
    "tools = [search_keyword]\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# 도구와 LLM 결합\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "\n",
    "########## 3. 노드 추가 ##########\n",
    "# 챗봇 함수 정의\n",
    "def chatbot(state: State):\n",
    "    # 메시지 호출 및 반환\n",
    "    return {\n",
    "        \"messages\": [llm_with_tools.invoke(state[\"messages\"])],\n",
    "        \"dummy_data\": \"[chatbot] 호출, dummy data\",  # 테스트를 위하여 더미 데이터를 추가합니다.\n",
    "    }\n",
    "\n",
    "\n",
    "# 상태 그래프 생성\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 챗봇 노드 추가\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "\n",
    "# 도구 노드 생성 및 추가\n",
    "tool_node = ToolNode(tools=tools)\n",
    "\n",
    "# 도구 노드 추가\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "# 조건부 엣지\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"chatbot\",\n",
    "    tools_condition,\n",
    ")\n",
    "\n",
    "########## 4. 엣지 추가 ##########\n",
    "\n",
    "# tools > chatbot\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "\n",
    "# START > chatbot\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# chatbot > END\n",
    "graph_builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "########## 5. 그래프 컴파일 ##########\n",
    "\n",
    "# 그래프 빌더 컴파일\n",
    "graph = graph_builder.compile()\n",
    "\n",
    "########## 6. 그래프 시각화 ##########\n",
    "# 그래프 시각화\n",
    "visualize_graph(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea7dc51",
   "metadata": {},
   "source": [
    "## StateGraph의 `stream` 메서드\n",
    "\n",
    "`stream` 메서드는 단일 입력에 대한 그래프 단계를 스트리밍하는 기능을 제공합니다.\n",
    "\n",
    "**매개변수**\n",
    "- `input` (Union[dict[str, Any], Any]): 그래프에 대한 입력\n",
    "- `config` (Optional[RunnableConfig]): 실행 구성\n",
    "- `stream_mode` (Optional[Union[StreamMode, list[StreamMode]]]): 출력 스트리밍 모드\n",
    "- `output_keys` (Optional[Union[str, Sequence[str]]]): 스트리밍할 키\n",
    "- `interrupt_before` (Optional[Union[All, Sequence[str]]]): 실행 전에 중단할 노드\n",
    "- `interrupt_after` (Optional[Union[All, Sequence[str]]]): 실행 후에 중단할 노드\n",
    "- `debug` (Optional[bool]): 디버그 정보 출력 여부\n",
    "- `subgraphs` (bool): 하위 그래프 스트리밍 여부\n",
    "\n",
    "**반환값**\n",
    "- Iterator[Union[dict[str, Any], Any]]: 그래프의 각 단계 출력. 출력 형태는 `stream_mode`에 따라 다름\n",
    "\n",
    "**주요 기능**\n",
    "1. 입력된 설정에 따라 그래프 실행을 스트리밍 방식으로 처리\n",
    "2. 다양한 스트리밍 모드 지원 (`values`, `updates`, `debug`)\n",
    "3. 콜백 관리 및 오류 처리\n",
    "4. 재귀 제한 및 중단 조건 처리\n",
    "\n",
    "**스트리밍 모드**\n",
    "- `values`: 각 단계의 현재 상태 값 출력\n",
    "- `updates`: 각 단계의 상태 업데이트만 출력\n",
    "- `debug`: 각 단계의 디버그 이벤트 출력"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3ae917b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 상태를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-tag\"],  # Tag\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e931823b",
   "metadata": {},
   "source": [
    "`config` 를 설정하고 스트리밍 출력을 진행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab5c0bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "for event in graph.stream(input=input, config=config):\n",
    "    for key, value in event.items():\n",
    "        print(f\"\\n[ {key} ]\\n\")\n",
    "        # value 에 messages 가 존재하는 경우\n",
    "        if \"messages\" in value:\n",
    "            messages = value[\"messages\"]\n",
    "            # 가장 최근 메시지 1개만 출력합니다.\n",
    "            value[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0d00b51",
   "metadata": {},
   "source": [
    "### `output_keys` 옵션\n",
    "\n",
    "`output_keys` 옵션은 스트리밍할 키를 지정하는 데 사용됩니다.\n",
    "\n",
    "list 형식으로 지정할 수 있으며, **channels 에 정의된 키 중 하나** 여야 합니다.\n",
    "\n",
    "**팁(tip)**\n",
    "\n",
    "- 매 단계마다 출력되는 State key 가 많은 경우, 일부만 스트리밍하고 싶은 경우에 유용합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe63155",
   "metadata": {},
   "outputs": [],
   "source": [
    "# channels 에 정의된 키 목록을 출력합니다.\n",
    "print(list(graph.channels.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87efc73e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 State 를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-rag\"],  # Tag\n",
    ")\n",
    "\n",
    "for event in graph.stream(\n",
    "    input=input,\n",
    "    config=config,\n",
    "    output_keys=[\"dummy_data\"],  # messages 를 추가해 보세요!\n",
    "):\n",
    "    for key, value in event.items():\n",
    "        # key 는 노드 이름\n",
    "        print(f\"\\n[ {key} ]\\n\")\n",
    "\n",
    "        # dummy_data 가 존재하는 경우\n",
    "        if value:\n",
    "            # value 는 노드의 출력값\n",
    "            print(value.keys())\n",
    "            # dummy_data key 가 존재하는 경우\n",
    "            if \"dummy_data\" in value:\n",
    "                print(value[\"dummy_data\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cecf357",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 State 를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-rag\"],  # Tag\n",
    ")\n",
    "\n",
    "for event in graph.stream(\n",
    "    input=input,\n",
    "    config=config,\n",
    "    output_keys=[\"messages\"],  # messages 만 출력\n",
    "):\n",
    "    for key, value in event.items():\n",
    "        # messages 가 존재하는 경우\n",
    "        if value and \"messages\" in value:\n",
    "            # key 는 노드 이름\n",
    "            print(f\"\\n[ {key} ]\\n\")\n",
    "            # messages 의 마지막 요소의 content 를 출력합니다.\n",
    "            print(value[\"messages\"][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f6be39d",
   "metadata": {},
   "source": [
    "### `stream_mode` 옵션\n",
    "\n",
    "`stream_mode` 옵션은 스트리밍 출력 모드를 지정하는 데 사용됩니다.\n",
    "\n",
    "- `values`: 각 단계의 현재 상태 값 출력 \n",
    "- `updates`: 각 단계의 상태 업데이트만 출력 (기본값)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f700fc",
   "metadata": {},
   "source": [
    "#### `stream_mode = \"values\"`\n",
    "\n",
    "`values` 모드는 각 단계의 현재 상태 값을 출력합니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "`event.items()`\n",
    "\n",
    "- `key`: State 의 key 값\n",
    "- `value`: State 의 key 에 대한하는 value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89eff21c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 State 를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-rag\"],  # Tag\n",
    ")\n",
    "\n",
    "# values 모드로 스트리밍 출력\n",
    "for event in graph.stream(\n",
    "    input=input,\n",
    "    stream_mode=\"values\",  # 기본값\n",
    "):\n",
    "    for key, value in event.items():\n",
    "        # key 는 state 의 key 값\n",
    "        print(f\"\\n[ {key} ]\\n\")\n",
    "        if key == \"messages\":\n",
    "            print(f\"메시지 개수: {len(value)}\")\n",
    "            # print(value)\n",
    "    print(\"===\" * 10, \" 단계 \", \"===\" * 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef295683",
   "metadata": {},
   "source": [
    "#### `stream_mode = \"updates\"`\n",
    "\n",
    "`updates` 모드는 각 단계에 대한 업데이트된 State 만 내보냅니다. \n",
    "\n",
    "- 출력은 노드 이름을 key 로, 업데이트된 값을 values 으로 하는 `dictionary` 입니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "`event.items()`\n",
    "\n",
    "- `key`: 노드(Node) 의 이름\n",
    "- `value`: 해당 노드(Node) 단계에서의 출력 값(dictionary). 즉, 여러 개의 key-value 쌍을 가진 dictionary 입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2124b916",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 State 를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-rag\"],  # Tag\n",
    ")\n",
    "\n",
    "# updates 모드로 스트리밍 출력\n",
    "for event in graph.stream(\n",
    "    input=input,\n",
    "    stream_mode=\"updates\",  # 기본값\n",
    "):\n",
    "    for key, value in event.items():\n",
    "        # key 는 노드 이름\n",
    "        print(f\"\\n[ {key} ]\\n\")\n",
    "\n",
    "        # value 는 노드의 출력값\n",
    "        print(value.keys())\n",
    "\n",
    "        # value 에는 state 가 dict 형태로 저장(values 의 key 값)\n",
    "        if \"messages\" in value:\n",
    "            print(f\"메시지 개수: {len(value['messages'])}\")\n",
    "            # print(value[\"messages\"])\n",
    "    print(\"===\" * 10, \" 단계 \", \"===\" * 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d98b21a8",
   "metadata": {},
   "source": [
    "### `interrupt_before` 와 `interrupt_after` 옵션\n",
    "\n",
    "`interrupt_before` 와 `interrupt_after` 옵션은 스트리밍 중단 시점을 지정하는 데 사용됩니다.\n",
    "\n",
    "- `interrupt_before`: 지정된 노드 이전에 스트리밍 중단\n",
    "- `interrupt_after`: 지정된 노드 이후에 스트리밍 중단\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f583a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 State 를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-rag\"],  # Tag\n",
    ")\n",
    "\n",
    "for event in graph.stream(\n",
    "    input=input,\n",
    "    config=config,\n",
    "    stream_mode=\"updates\",  # 기본값\n",
    "    interrupt_before=[\"tools\"],  # tools 노드 이전에 스트리밍 중단\n",
    "):\n",
    "    for key, value in event.items():\n",
    "        # key 는 노드 이름\n",
    "        print(f\"\\n[{key}]\\n\")\n",
    "\n",
    "        # value 는 노드의 출력값\n",
    "        if isinstance(value, dict):\n",
    "            print(value.keys())\n",
    "            if \"messages\" in value:\n",
    "                print(value[\"messages\"])\n",
    "\n",
    "        # value 에는 state 가 dict 형태로 저장(values 의 key 값)\n",
    "        if \"messages\" in value:\n",
    "            print(f\"메시지 개수: {len(value['messages'])}\")\n",
    "    print(\"===\" * 10, \" 단계 \", \"===\" * 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33dc67a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문\n",
    "question = \"2024년 노벨 문학상 관련 뉴스를 알려주세요.\"\n",
    "\n",
    "# 초기 입력 State 를 정의\n",
    "input = State(dummy_data=\"테스트 문자열\", messages=[(\"user\", question)])\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    "    tags=[\"my-rag\"],  # Tag\n",
    ")\n",
    "\n",
    "for event in graph.stream(\n",
    "    input=input,\n",
    "    config=config,\n",
    "    stream_mode=\"updates\",\n",
    "    interrupt_after=[\"tools\"],  # tools 실행 후 interrupt\n",
    "):\n",
    "    for value in event.values():\n",
    "        # key 는 노드 이름\n",
    "        print(f\"\\n[{key}]\\n\")\n",
    "\n",
    "        if isinstance(value, dict):\n",
    "            # value 는 노드의 출력값\n",
    "            print(value.keys())\n",
    "            if \"messages\" in value:\n",
    "                print(value[\"messages\"])\n",
    "\n",
    "        # value 에는 state 가 dict 형태로 저장(values 의 key 값)\n",
    "        if \"messages\" in value:\n",
    "            print(f\"메시지 개수: {len(value['messages'])}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
